package com.example.llm.sgg.sgglangchain.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeServerlessIndexConfig;
import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Arrays;
import java.util.List;

@Configuration
public class XiaozhiAgentConfig {

    @Autowired
    private XiaozhiAgentChatMemoryStore xiaozhiAgentChatMemoryStore;

    @Bean(value = "xiaozhiChatMemoryProvider")
    public ChatMemoryProvider xiaozhiChatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(30) // 由于调用工具方法、以及获取执行方法的返回结果都占用 会话 上下文的大小，所以此处稍微调大些
//                .chatMemoryStore(new InMemoryChatMemoryStore())
                .chatMemoryStore(xiaozhiAgentChatMemoryStore)
                .build();
    }

    @Bean
    public ContentRetriever contentRetrieverXiaozhi() {
        // 使用FileSystemDocumentLoader读取指定目录下的知识库文档
        // 并使用默认的文档解析器对文档进行解析
//        Document document1 = FileSystemDocumentLoader.loadDocument("E:/knowledge/医院信息.md");
//        Document document2 = FileSystemDocumentLoader.loadDocument("E:/knowledge/科室信息.md");
//        Document document3 = FileSystemDocumentLoader.loadDocument("E:/knowledge/神经内科.md");
        Document document1 = Document.from("医院名称：\n" +
                "北京协和医院\n" +
                "\n" +
                "医院地址：\n" +
                "（东单院区）北京市东城区帅府园一号，邮编：100730\n" +
                "（西单院区）北京市西城区大木仓胡同41号，邮编：100032\n" +
                "\n" +
                "门诊开放时间：\n" +
                "工作日：8:00 - 17:30\n" +
                "急诊：24小时开放\n" +
                "\n" +
                "服务热线：\n" +
                "（东单院区）010-69151188\n" +
                "（西单院区）010-69158100\n" +
                "\n" +
                "医院简介：\n" +
                "北京协和医院是集医疗、教学、科研于一体的现代化综合三级甲等医院，是国家卫生健康委指定的全国疑难重症诊治指导中心，最早承担外宾医疗任务的医院之一，也是高等医学教育和住院医师规范化培训国家级示范基地，临床医学研究和技术创新的国家级核心基地。以学科齐全、技术力量雄厚、特色专科突出、多学科综合优势强大享誉海内外。在国家三级公立医院绩效考核中排名第一。\n" +
                "\n" +
                "东单院区乘车路线：\n" +
                "1. 106, 108, 110, 111, 116路到东单路口北\n" +
                "2. 41, 104快路到东单路口南");
        Document document2 = Document.from("北京协和医院科室详细介绍：\n" +
                "\n" +
                "1、内科学系：心内科、呼吸与危重症医学科、消化内科、肾内科、血液内科、风湿免疫科、感染内科、全科医学科（普通内科）、肿瘤内科、MIGU、老年医学科、内分泌科、重症医学科（ICU）、儿科、神经内科、心理医学科、皮肤科、变态（过敏）反应科、急诊科、中医科、康复医学科、临床营养科、罕见病医学科。\n" +
                "\n" +
                "2、外科学系：基本外科、骨科、心外科、胸外科、泌尿外科、神经外科、血管外科、整形美容外科、乳腺外科、肝脏外科、麻醉科。\n" +
                "\n" +
                "3、其他科室：妇产科、眼科、耳鼻喉科、口腔科、超声医学科、病理科、检验科、放射科、放射治疗科、核医学科、输血科、药剂科、病案科、国际医疗部、健康医学部（体检中心）、临床药理研究中心。");
        Document document3 = Document.from("北京协和医院神经内科详细介绍：\n" +
                "\n" +
                "目前床位数为200张（含4个病区和神经科重症监护病房），年门诊量15万余人次，年住院患者数量1万余人次。目前亚学科发展齐全，包括：脑血管病、神经介入和神经血管内治疗、帕金森病和运动障碍、神经康复、神经重症、神经遗传性疾病、神经心理和痴呆、癫痫和睡眠障碍、神经肌肉病以及神经免疫病等。科室设有脑血管病诊治中心，已开展急性缺血性脑血管病动、静脉溶栓和神经血管内介入治疗，脑血流灌注/脑血管动力学检查，动脉粥样硬化模块及微粒子监测，抗血小板药物等相关基因多态性检测等；建有神经肌肉病理和分子生物学检查实验室，已开展神经、肌肉活检和组织化学染色以及进行常见神经疾病的基因诊断；有独立的神经科康复训练中心，瘫痪患者有上下肢机器人供功能锻炼康复，肌张力障碍患者可行肉毒毒素治疗；有神经电生理中心，可进行视频脑电图检测，睡眠脑电监测以及常规的神经肌电图检查。\n" +
                "\n" +
                "神经内科创建于上世纪70年代，本学科人才结构合理，拥有一批丰富力强、富有创新能力的技术骨干。全科在职职工近150人，医生42人，护士95人，技术员10人，其中正高14人，副高13人，主管护师7人；博士生导师5人，硕士生导师19人；年青医生都具有博士学位。\n" +
                "\n" +
                "科研教学水平在国内处于领先地位，是国家首批院士、博士后培养点及住院医师培训基地。每年培养硕士、博士20余人，接受进修医生30余人。目前共承担国家和省部级项目70余项，获得省级科技进步奖10余项，国家科技进步二等奖1项-童尊增教授等主持的《低血钾软病的研究》，为成功制服曾肆虐10余篇临床医学的低血钾软病作出了重大贡献。到2014年底，我科累计发表论文1400余篇，SCI收录140余篇，其中包括PNAS，J Neurosci，Gila，J Alzheimers Dis等国际知名杂志发表有影响力的学术论文。\n" +
                "\n" +
                "我科是国家临床重点专科，中华医学会湖北省神经病学分会和中华预防医学会湖北省中预防与控制专业委员会主任委员单位，卫生部神经内科医疗质量控制中心湖北省分中心暨湖北省神经内科医疗质量控制中心所在单位。\n" +
                "\n" +
                "科室领导：\n" +
                "主任：胡波（负责科室各项工作的统筹规划）\n" +
                "副主任：王涛（负责科室科研工作）\n" +
                "副主任：曹非（负责科室教学工作）\n" +
                "副主任：李红龙（负责科室医疗工作）");
        List<Document> documents = Arrays.asList(document1, document2, document3);

        // 使用内存向量存储
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        //使用默认的文档分割器
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

        // 从嵌入存储（EmbeddingStore）里检索和查询内容相关的信息
        // TODO 如何将用户动态的参数传进来，比如过滤条件
        return EmbeddingStoreContentRetriever.builder()
                .displayName("内存向量数据库")
                .embeddingStore(embeddingStore)
                .maxResults(1)
                .minScore(0.8)
//                .dynamicFilter()
                .build();
    }

    @Autowired
    private EmbeddingModel embeddingModel;

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        // 创建向量存储
        EmbeddingStore<TextSegment> embeddingStore = PineconeEmbeddingStore.builder()
                .apiKey("pcsk_6Gt8de_NtkfpxVYZp9s3yesgWandLaH3kNZkfnAAKwThqnGZYVKb5DuANckVFsqNnPQHJW")
                .index("xiaozhi-index") // 如果指定的索引不存在，将创建一个新的索引
                .nameSpace("xiaozhi-namespace") // 如果指定的名称空间不存在，将创建一个新的名称空间
                .createIndex(PineconeServerlessIndexConfig.builder()
                        .cloud("AWS") // 指定索引部署在 AWS 云服务上。
                        .region("us-east-1") // 指定索引所在的 AWS 区域为 us-east-1。
                        .dimension(embeddingModel.dimension()) // 指定索引的向量维度，该维度与 embeddedModel 生成的向量维度相同。
                        .build())
                .build();
        return embeddingStore;
    }

}
